import os
import sys
import subprocess
import time
import threading
import wave
import pyaudio
import numpy as np
from datetime import datetime
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QPushButton, QLabel, QRadioButton,
                           QButtonGroup, QMessageBox, QFileDialog, QDialog,
                           QSpinBox, QFormLayout, QCheckBox, QDoubleSpinBox)
from PyQt5.QtCore import Qt, QTimer, QSize, pyqtSignal, QPoint
from PyQt5.QtGui import QIcon, QPainter, QPen, QColor

def resource_path(relative_path):
    """ 获取资源的绝对路径 """
    try:
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.abspath(".")
    return os.path.join(base_path, relative_path)

class CountdownOverlay(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setAttribute(Qt.WA_ShowWithoutActivating)
        self.setFixedSize(150, 150)
        
        self.parent_center = parent.geometry().center()
        self.move(self.parent_center - QPoint(75, 75))
        
        self.countdown = 3
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_countdown)
        self.last_update_time = time.time()
        
    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.setBrush(QColor(0, 0, 0, 150))
        painter.setPen(Qt.NoPen)
        painter.drawRoundedRect(self.rect(), 10, 10)
        
        painter.setPen(QColor(255, 255, 255))
        if self.countdown > 0:
            text = str(self.countdown)
            font = painter.font()
            font.setPointSize(60)
            font.setBold(True)
            painter.setFont(font)
            painter.drawText(self.rect(), Qt.AlignCenter, text)
        
    def start_countdown(self):
        self.show()
        self.last_update_time = time.time()
        self.timer.start(100)
        
    def update_countdown(self):
        current_time = time.time()
        if current_time - self.last_update_time >= 1:
            self.countdown -= 1
            self.last_update_time = current_time
            self.update()
            
            if self.countdown <= 0:
                self.timer.stop()
                self.hide()
                self.countdown_finished.emit()
    
    countdown_finished = pyqtSignal()

class SettingsDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.parent = parent
        self.setWindowTitle("设置")
        self.setWindowIcon(QIcon(resource_path("resources/settings.png")))
        self.setFixedSize(420, 280)
        self.setWindowFlags(self.windowFlags() & ~Qt.WindowContextHelpButtonHint)  # 移除问号按钮
        
        self.init_ui()
        
    def init_ui(self):
        layout = QVBoxLayout()
        layout.setContentsMargins(15, 15, 15, 15)
        layout.setSpacing(15)
        
        form_layout = QFormLayout()
        form_layout.setVerticalSpacing(12)
        form_layout.setHorizontalSpacing(15)
        
        # 输出目录设置
        self.dir_display = QLabel(self.parent.output_dir)
        self.dir_display.setStyleSheet("color: #CCCCCC;")
        self.dir_btn = QPushButton("更改")
        self.dir_btn.setFixedWidth(80)
        
        dir_layout = QHBoxLayout()
        dir_layout.addWidget(self.dir_display)
        dir_layout.addWidget(self.dir_btn)
        dir_layout.setSpacing(10)
        form_layout.addRow("输出目录:", dir_layout)
        
        # 帧率设置
        self.framerate = QSpinBox()
        self.framerate.setRange(10, 60)
        self.framerate.setValue(self.parent.framerate)
        self.framerate.setFixedWidth(120)
        self.framerate.setStyleSheet("QSpinBox::up-button { width: 20px; } QSpinBox::down-button { width: 20px; }")
        
        fps_label = QLabel("fps (10-60，值越大越流畅)")
        fps_label.setStyleSheet("color: #AAAAAA; font-size: 11px;")
        
        fps_layout = QHBoxLayout()
        fps_layout.addWidget(self.framerate)
        fps_layout.addWidget(fps_label)
        fps_layout.setSpacing(5)
        form_layout.addRow("帧率:", fps_layout)
        
        # 音频设置
        self.audio_checkbox = QCheckBox("录制系统音频")
        self.audio_checkbox.setChecked(self.parent.record_audio)
        form_layout.addRow("音频:", self.audio_checkbox)
        
        # 音频增益设置
        self.audio_gain = QDoubleSpinBox()
        self.audio_gain.setRange(1.0, 20.0)
        self.audio_gain.setSingleStep(0.5)
        self.audio_gain.setValue(self.parent.audio_gain)
        self.audio_gain.setFixedWidth(120)
        self.audio_gain.setStyleSheet("QDoubleSpinBox::up-button { width: 20px; } QDoubleSpinBox::down-button { width: 20px; }")
        
        gain_label = QLabel("x (1.0-20.0，值越大声音越大)")
        gain_label.setStyleSheet("color: #AAAAAA; font-size: 11px;")
        
        gain_layout = QHBoxLayout()
        gain_layout.addWidget(self.audio_gain)
        gain_layout.addWidget(gain_label)
        gain_layout.setSpacing(5)
        form_layout.addRow("音频增益:", gain_layout)
        
        layout.addLayout(form_layout)
        
        # 确认按钮
        btn_layout = QHBoxLayout()
        btn_layout.addStretch()
        
        self.cancel_btn = QPushButton("取消")
        self.cancel_btn.setFixedWidth(90)
        self.cancel_btn.clicked.connect(self.reject)
        btn_layout.addWidget(self.cancel_btn)
        
        self.ok_btn = QPushButton("确定")
        self.ok_btn.setFixedWidth(90)
        self.ok_btn.clicked.connect(self.accept)
        btn_layout.addWidget(self.ok_btn)
        
        layout.addLayout(btn_layout)
        self.setLayout(layout)
        
    def select_output_dir(self):
        dir = QFileDialog.getExistingDirectory(self, "选择目录", self.parent.output_dir)
        if dir:
            self.dir_display.setText(dir)
            self.dir_display.setToolTip(dir)

class RegionSelector(QWidget):
    region_selected = pyqtSignal(tuple)
    
    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.start_pos = self.end_pos = None

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        painter.fillRect(self.rect(), QColor(0, 0, 0, 100))
        
        if self.start_pos and self.end_pos:
            x, y = min(self.start_pos.x(), self.end_pos.x()), min(self.start_pos.y(), self.end_pos.y())
            w, h = abs(self.start_pos.x() - self.end_pos.x()), abs(self.start_pos.y() - self.end_pos.y())
            
            if w % 2 != 0: w += 1
            if h % 2 != 0: h += 1
            
            painter.setCompositionMode(QPainter.CompositionMode_Clear)
            painter.fillRect(x, y, w, h, Qt.transparent)
            
            painter.setCompositionMode(QPainter.CompositionMode_SourceOver)
            painter.setPen(QPen(Qt.red, 2, Qt.DashLine))
            painter.drawRect(x, y, w, h)
            painter.drawText(x, y - 10, f"{w}x{h}")

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.start_pos = self.end_pos = event.pos()

    def mouseMoveEvent(self, event):
        if self.start_pos:
            self.end_pos = event.pos()
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.start_pos:
            x1, y1 = self.start_pos.x(), self.start_pos.y()
            x2, y2 = event.pos().x(), event.pos().y()
            w, h = abs(x1-x2), abs(y1-y2)
            
            if w % 2 != 0: w += 1
            if h % 2 != 0: h += 1
                
            self.region_selected.emit((min(x1,x2), min(y1,y2), w, h))
            self.close()

class AudioRecorder:
    def __init__(self, gain=5.0):
        self.audio_frames = []
        self.recording = False
        self.audio_thread = None
        self.stream = None
        self.p = pyaudio.PyAudio()
        self.gain = gain
        
    def start_recording(self):
        self.recording = True
        self.audio_frames = []
        
        def audio_thread_func():
            self.stream = self.p.open(
                format=pyaudio.paInt16,
                channels=2,
                rate=44100,
                input=True,
                frames_per_buffer=1024,
                stream_callback=self.audio_callback
            )
            while self.recording:
                time.sleep(0.1)
            self.stream.stop_stream()
            self.stream.close()
            
        self.audio_thread = threading.Thread(target=audio_thread_func)
        self.audio_thread.start()
        
    def audio_callback(self, in_data, frame_count, time_info, status):
        if self.recording:
            audio_data = np.frombuffer(in_data, dtype=np.int16)
            processed_data = (audio_data * self.gain).clip(-32768, 32767).astype(np.int16)
            self.audio_frames.append(processed_data.tobytes())
        return (None, pyaudio.paContinue)
        
    def stop_recording(self):
        self.recording = False
        if self.audio_thread:
            self.audio_thread.join()
        return self.audio_frames
        
    def save_audio(self, filename, frames):
        wf = wave.open(filename, 'wb')
        wf.setnchannels(2)
        wf.setsampwidth(self.p.get_sample_size(pyaudio.paInt16))
        wf.setframerate(44100)
        wf.writeframes(b''.join(frames))
        wf.close()

class ScreenRecorder(QMainWindow):
    def __init__(self):
        super().__init__()
        self.recording = False
        self.ffmpeg_process = None
        self.audio_recorder = None
        self.region = None
        self.region_selector = None
        self.output_dir = os.path.join(os.getcwd(), "output")
        self.log_dir = os.path.join(os.getcwd(), "log")
        self.temp_dir = os.path.join(os.getcwd(), "temp")
        self.framerate = 30
        self.record_audio = True
        self.audio_gain = 5.0
        
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.temp_dir, exist_ok=True)
        
        self.init_ui()
        self.setWindowTitle("屏幕录制工具")
        self.setWindowIcon(QIcon(resource_path("resources/record.png")))
        self.setFixedSize(450, 200)
        self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)
        self.show()

    def init_ui(self):
        self.setStyleSheet("""
            QMainWindow, QDialog { background-color: #2D2D2D; font-family: "Microsoft YaHei"; }
            QRadioButton { color: #EEEEEE; font-size: 13px; spacing: 5px; }
            QRadioButton::indicator { width: 16px; height: 16px; }
            QRadioButton::indicator::unchecked { border: 2px solid #888888; border-radius: 8px; background: #444444; }
            QRadioButton::indicator::checked { border: 2px solid #FFA500; border-radius: 8px; background: #FFA500; }
            QLabel { color: #FFFFFF; font-size: 12px; }
            QPushButton { background-color: #3A3A3A; color: white; border: none; padding: 5px; min-width: 60px; }
            QPushButton:hover { background-color: #4A4A4A; }
            #timeLabel { font-size: 24px; color: #FFA500; }
            #statusLabel { font-size: 12px; color: #AAAAAA; }
            QSpinBox { background-color: #3A3A3A; color: white; border: 1px solid #555555; padding: 2px; }
            QCheckBox { color: #EEEEEE; font-size: 13px; spacing: 5px; }
            QCheckBox::indicator { width: 16px; height: 16px; }
            QCheckBox::indicator::unchecked { border: 2px solid #888888; background: #444444; }
            QCheckBox::indicator::checked { border: 2px solid #FFA500; background: #FFA500; }
            QDoubleSpinBox { background-color: #3A3A3A; color: white; border: 1px solid #555555; padding: 2px; }
        """)

        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QHBoxLayout()
        main_widget.setLayout(main_layout)
        
        # 左侧控制区
        control_panel = QWidget()
        control_layout = QVBoxLayout()
        control_panel.setLayout(control_layout)
        
        mode_group = QButtonGroup(self)
        self.fullscreen_radio = QRadioButton("全屏录制")
        self.region_radio = QRadioButton("区域录制")
        mode_group.addButton(self.fullscreen_radio)
        mode_group.addButton(self.region_radio)
        self.fullscreen_radio.setChecked(True)
        control_layout.addWidget(self.fullscreen_radio)
        control_layout.addWidget(self.region_radio)
        control_layout.addSpacing(10)
        
        btn_layout = QHBoxLayout()
        self.record_btn = QPushButton("开始")
        self.record_btn.setIcon(QIcon(resource_path("resources/record.png")))
        self.record_btn.clicked.connect(self.toggle_recording)
        btn_layout.addWidget(self.record_btn)
        
        self.open_btn = QPushButton("打开")
        self.open_btn.setIcon(QIcon(resource_path("resources/folder.png")))
        self.open_btn.clicked.connect(self.open_output_folder)
        btn_layout.addWidget(self.open_btn)
        
        self.settings_btn = QPushButton("设置")
        self.settings_btn.setIcon(QIcon(resource_path("resources/settings.png")))
        self.settings_btn.clicked.connect(self.show_settings)
        btn_layout.addWidget(self.settings_btn)
        
        control_layout.addLayout(btn_layout)
        
        # 右侧状态区
        status_panel = QWidget()
        status_layout = QVBoxLayout()
        status_panel.setLayout(status_layout)
        
        self.time_label = QLabel("00:00:00")
        self.time_label.setObjectName("timeLabel")
        status_layout.addWidget(self.time_label)
        
        self.status_label = QLabel("准备就绪")
        self.status_label.setObjectName("statusLabel")
        status_layout.addWidget(self.status_label)
        
        main_layout.addWidget(control_panel)
        main_layout.addWidget(status_panel)
        
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_timer)
        self.countdown_overlay = CountdownOverlay(self)
        self.countdown_overlay.countdown_finished.connect(self.start_recording)
        self.start_time = 0

    def toggle_recording(self):
        if self.recording:
            self.stop_recording()
        else:
            if self.region_radio.isChecked():
                self.select_region()
            else:
                self.start_countdown()

    def select_region(self):
        self.hide()
        self.region_selector = RegionSelector()
        self.region_selector.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.region_selector.region_selected.connect(self.on_region_selected)
        self.region_selector.showFullScreen()

    def on_region_selected(self, region):
        x, y, w, h = region
        if w % 2 != 0: w -= 1
        if h % 2 != 0: h -= 1
        
        self.region = (x, y, w, h)
        self.show()
        reply = QMessageBox.question(
            self, '确认区域', f"区域大小: {w}x{h}\n开始录制?",
            QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if reply == QMessageBox.Yes:
            self.start_countdown()

    def start_countdown(self):
        self.status_label.setText("准备录制...")
        self.countdown_overlay.countdown = 3
        self.countdown_overlay.move(self.geometry().center() - QPoint(75, 75))
        self.countdown_overlay.start_countdown()

    def start_recording(self):
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.video_file = os.path.join(self.temp_dir, f"video_{timestamp}.mp4")
            self.audio_file = os.path.join(self.temp_dir, f"audio_{timestamp}.wav")
            self.output_file = os.path.join(self.output_dir, f"recording_{timestamp}.mp4")
            log_file = os.path.join(self.log_dir, f"ffmpeg_{timestamp}.log")

            # 启动音频录制
            if self.record_audio:
                try:
                    self.audio_recorder = AudioRecorder(gain=self.audio_gain)
                    self.audio_recorder.start_recording()
                except Exception as e:
                    self.record_audio = False
                    self.status_label.setText(f"音频启动失败: {str(e)}")

            # 构建视频录制命令
            cmd = [
                resource_path("resources/ffmpeg.exe"),
                '-y',
                '-f', 'gdigrab',
                '-framerate', str(self.framerate),
                '-probesize', '50M',
                '-draw_mouse', '1',
            ]

            if not self.fullscreen_radio.isChecked():
                x, y, w, h = self.region
                cmd.extend([
                    '-offset_x', str(x),
                    '-offset_y', str(y),
                    '-video_size', f"{w}x{h}"
                ])

            cmd.extend([
                '-i', 'desktop',
                '-c:v', 'libx264',
                '-preset', 'fast',
                '-crf', '18',
                '-pix_fmt', 'yuv420p',
                '-movflags', '+faststart',
                self.video_file
            ])

            print("执行命令:", " ".join(cmd))
            
            with open(log_file, 'w') as f:
                self.ffmpeg_process = subprocess.Popen(
                    cmd, stdin=subprocess.PIPE, 
                    stdout=f, stderr=subprocess.STDOUT,
                    creationflags=subprocess.CREATE_NO_WINDOW)
            
            self.recording = True
            self.start_time = time.time()
            self.timer.start(1000)
            self.status_label.setText("录制中...")
            self.record_btn.setText("停止")
            self.record_btn.setIcon(QIcon(resource_path("resources/stop.png")))

        except Exception as e:
            self.status_label.setText(f"启动失败: {str(e)}")
            QMessageBox.critical(self, "错误", f"启动失败: {str(e)}")

    def stop_recording(self):
        if not self.recording:
            return

        # 停止视频录制
        if self.ffmpeg_process:
            try:
                self.ffmpeg_process.stdin.write(b'q\n')
                self.ffmpeg_process.stdin.flush()
                self.ffmpeg_process.wait(timeout=3)
            except:
                self.ffmpeg_process.terminate()
            finally:
                self.ffmpeg_process = None

        # 停止音频录制并保存
        audio_frames = None
        if self.record_audio and self.audio_recorder:
            try:
                audio_frames = self.audio_recorder.stop_recording()
                if audio_frames:
                    self.audio_recorder.save_audio(self.audio_file, audio_frames)
            except Exception as e:
                self.status_label.setText(f"音频保存失败: {str(e)}")
                self.record_audio = False

        self.recording = False
        self.timer.stop()
        
        # 合并音视频
        if self.record_audio and audio_frames and os.path.exists(self.video_file) and os.path.exists(self.audio_file):
            self.merge_audio_video()
        elif os.path.exists(self.video_file):
            os.rename(self.video_file, self.output_file)
            self.status_label.setText("视频已保存 (无音频)")
        else:
            self.status_label.setText("录制失败，请检查日志")

        self.record_btn.setText("开始")
        self.record_btn.setIcon(QIcon(resource_path("resources/record.png")))

    def merge_audio_video(self):
        try:
            cmd = [
                resource_path("resources/ffmpeg.exe"),
                '-y',
                '-i', self.video_file,
                '-i', self.audio_file,
                '-c:v', 'copy',
                '-c:a', 'aac',
                '-strict', 'experimental',
                '-map', '0:v:0',
                '-map', '1:a:0',
                '-shortest',
                self.output_file
            ]
            
            subprocess.run(cmd, check=True, creationflags=subprocess.CREATE_NO_WINDOW)
            self.status_label.setText("音视频合并完成")
            
            # 清理临时文件
            try:
                os.remove(self.video_file)
                os.remove(self.audio_file)
            except:
                pass
                
        except Exception as e:
            self.status_label.setText(f"合并失败: {str(e)}")
            # 保留原始视频文件
            if os.path.exists(self.video_file):
                os.rename(self.video_file, self.output_file)

    def open_output_folder(self):
        if os.path.exists(self.output_dir):
            os.startfile(self.output_dir)

    def update_timer(self):
        elapsed = int(time.time() - self.start_time)
        self.time_label.setText(f"{elapsed//3600:02d}:{(elapsed%3600)//60:02d}:{elapsed%60:02d}")

    def show_settings(self):
        dialog = SettingsDialog(self)
        if dialog.exec_() == QDialog.Accepted:
            self.output_dir = dialog.dir_display.text()
            self.framerate = dialog.framerate.value()
            self.record_audio = dialog.audio_checkbox.isChecked()
            self.audio_gain = dialog.audio_gain.value()

    def closeEvent(self, event):
        if self.recording:
            reply = QMessageBox.question(
                self, '确认退出', "正在录制中，确定退出吗?",
                QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if reply == QMessageBox.Yes:
                self.stop_recording()
                event.accept()
            else:
                event.ignore()
        else:
            event.accept()

if __name__ == "__main__":
    os.environ['QT_LOGGING_RULES'] = 'qt.qpa.fonts=false'
    if hasattr(Qt, 'AA_EnableHighDpiScaling'):
        QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    
    app = QApplication(sys.argv)
    app.setStyle('Fusion')
    font = app.font()
    font.setFamily("Microsoft YaHei")
    app.setFont(font)
    
    window = ScreenRecorder()
    window.show()
    sys.exit(app.exec_())